#include <driver_functions.h>

#include "inc/core/ucgrid3d.h"
#include <stdlib.h>
#include <stdio.h>
#include "math.h"
#include "gsl/gsl_randist.h"
#include "inc/core/lib.h"
#include "cufft.h"

void empty_grid(ucgrid3d* grid)
{
    cudaExtent extent = make_cudaExtent(sizeof(cufftComplex)*grid->Nx, grid->Ny, grid->Nz);
    cudaMemset3D(grid->rho, 0, extent);
    cudaMemset3D(grid->phi, 0, extent);
}

ucgrid3d* ucgrid3d_init_grid(ucgrid3d* grid, int Nx, int Ny, int Nz,
                             floatdig dx, floatdig dy, floatdig dz)
{
    int i;
    int j;
    int k;
#ifdef VERBOSE
    printf("ucgrid3d_init_grid: init grid with sizes Nx: %d, Ny: %d, Nz: %d\n", Nx, Ny, Nz);
#endif
    grid->dx = dx;
    grid->dy = dy;
    grid->dz = dz;
    grid->Nx = Nx;
    grid->Ny = Ny;
    grid->Nz = Nz;
    cudaExtent extent = make_cudaExtent(sizeof(cufftComplex)*Nx, Ny, Nz);
    cudaMalloc3D (&(grid->rho), extent);
    cudaMalloc3D (&(grid->phi), extent);
    cudaMemset3D(grid->rho, 0, extent);
    cudaMemset3D(grid->phi, 0, extent);
    return grid;
}

void ucgrid3d_free_grid(ucgrid3d* grid)
{
    int Nx = grid->Nx;
    int Ny = grid->Ny;
    int i;
    int j;

#ifdef VERBOSE
    printf("ucgrid3d_free_grid: Free memory from grid's data\n");
#endif
    cudaFree(&(grid->rho));
    cudaFree(&(grid->phi));

}

//void ucgrid3d_get_cell_pos(ucgrid3d* grid, npos3d* n, pos3d* p)
//{
//    p->v[X] = n->v[X] * grid->dx;
//    p->v[Y] = n->v[Y] * grid->dy;
//    p->v[Z] = n->v[Z] * grid->dz;
//}
//
//void ucgrid3d_get_full_size_of_system(ucgrid3d* grid, pos3d* Ls)
//{
//    Ls->v[X] = grid->Nx * grid->dx;
//    Ls->v[Y] = grid->Ny * grid->dy;
//    Ls->v[Z] = grid->Nz * grid->dz;
//}
//
void ucgrid3d_get_neighbour_cells(ucgrid3d* grid, pos3d* ppos, npos3d* lc, npos3d* rc)
{
    floatdig x = ppos->v[X] / grid->dx;
    floatdig y = ppos->v[Y] / grid->dy;
    floatdig z = ppos->v[Z] / grid->dz;

    lc->v[X] = floor(x);
    rc->v[X] = ceil(x);
    lc->v[Y] = floor(y);
    rc->v[Y] = ceil(y);
    lc->v[Z] = floor(z);
    rc->v[Z] = ceil(z);
}
//

__device__ 
floatdig* get_element(cudaPitchedPtr* ptr, cudaExtent* extent, int sX, int sY, int sZ)
{
    floatdig* devPtr = (floatdig*)ptr->ptr; 
    size_t pitch = ptr->pitch; 
    size_t slicePitch = pitch * extent->height;
    floatdig* slice = devPtr + sZ * slicePitch;
    floatdig* row = (floatdig*)(slice + sY * pitch);
    return  (floatdig*)&row[sX];
}


__device__ void get_cell_pos(size3d* cellsize, npos3d* n, pos3d* p)
{
    p->v[X] = n->v[X] * cellsize->v[X];
    p->v[Y] = n->v[Y] * cellsize->v[Y];
    p->v[Z] = n->v[Z] * cellsize->v[Z];
}

__device__ void get_neighbour_cells(particle* p, npos3d* lc, npos3d* rc, size3d* cellsize)
{
    pos3d* ppos = &(p->pos);
    floatdig x = ppos->v[X] / cellsize->v[X];
    floatdig y = ppos->v[Y] / cellsize->v[Y];
    floatdig z = ppos->v[Z] / cellsize->v[Z];
    lc->v[X] = floor(x);
    rc->v[X] = ceil(x);
    lc->v[Y] = floor(y);
    rc->v[Y] = ceil(y);
    lc->v[Z] = floor(z);
    rc->v[Z] = ceil(z);
}


__global__ void do_weight_particles_rho(cudaPitchedPtr rho, cudaExtent extent, particle* pbunch, floatdig q, floatdig dx, floatdig dy, floatdig dz)
{
    int nparticle = blockIdx.x*blockDim.x + threadIdx.x;
    npos3d lc = {0};
    npos3d rc = {0};
    
    npos3d nlc = {0};
    npos3d nrc = {0};
    pos3d plc = {0};
    pos3d prc = {0};

    pos3d plll = {0};
    pos3d prll = {0};
    pos3d plrl = {0};
    pos3d plrr = {0};
    pos3d pllr = {0};
    pos3d prlr = {0};
    pos3d prrr = {0};
    pos3d prrl = {0};

    floatdig lllq = 0;
    floatdig rllq = 0;
    floatdig lrlq = 0;
    floatdig lrrq = 0;
    floatdig llrq = 0;
    floatdig rlrq = 0;
    floatdig rrrq = 0;
    floatdig rrlq = 0;
    floatdig* rhoc = NULL;
    
    floatdig dv = dx*dy*dz;
    floatdig qdv = q*dv;
    pos3d* ppos = &(pbunch[nparticle].pos);
    
    size3d cellsize;
    cellsize.v[X] = dx;
    cellsize.v[Y] = dy;
    cellsize.v[Z] = dz;	
    
    get_neighbour_cells(&(pbunch[nparticle]), &lc, &rc, &cellsize);
    get_cell_pos(&cellsize, &nlc, &plc);
    get_cell_pos(&cellsize, &nrc, &prc);
    plll.v[X] = plc.v[X];
    plll.v[Y] = plc.v[Y];
    plll.v[Z] = plc.v[Z];

    prll.v[X] = prc.v[X];
    prll.v[Y] = plc.v[Y];
    prll.v[Z] = plc.v[Z];

    plrl.v[X] = plc.v[X];
    plrl.v[Y] = prc.v[Y];
    plrl.v[Z] = plc.v[Z];

    prrl.v[X] = prc.v[X];
    prrl.v[Y] = prc.v[Y];
    prrl.v[Z] = plc.v[Z];

    prrr.v[X] = prc.v[X];
    prrr.v[Y] = prc.v[Y];
    prrr.v[Z] = prc.v[Z];

    prlr.v[X] = prc.v[X];
    prlr.v[Y] = plc.v[Y];
    prlr.v[Z] = prc.v[Z];

    pllr.v[X] = plc.v[X];
    pllr.v[Y] = plc.v[Y];
    pllr.v[Z] = prc.v[Z];

    plrr.v[X] = plc.v[X];
    plrr.v[Y] = prc.v[Y];
    plrr.v[Z] = prc.v[Z];
    
    lllq = qdv * (prll.v[X] - ppos->v[X])*(plrl.v[Y] - ppos->v[Y])*(pllr.v[Z] - ppos->v[Z]) / dv;
    rllq = qdv * (ppos->v[X] - plll.v[X])*(prrl.v[Y] - ppos->v[Y])*(prlr.v[Z] - ppos->v[Z]) / dv;

   lrlq = qdv * (prrl.v[X] - ppos->v[X])*(ppos->v[Y] - plll.v[Y])*(plrr.v[Z] - ppos->v[Z]) / dv;
   rrlq = qdv * (ppos->v[X] - plrl.v[X])*(ppos->v[Y] - prll.v[Y])*(prrr.v[Z] - ppos->v[Z]) / dv;

   rrrq = qdv * (ppos->v[X] - plrr.v[X])*(ppos->v[Y] - prlr.v[Y])*(ppos->v[Z] - prrl.v[Z]) / dv;
   rlrq = qdv * (ppos->v[X] - pllr.v[X])*(prrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - prll.v[Z]) / dv;

   llrq = qdv * (prlr.v[X] - ppos->v[X])*(plrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - plll.v[Z]) / dv;
   lrrq = qdv * (prrr.v[X] - ppos->v[X])*(ppos->v[Y] - pllr.v[Y])*(ppos->v[Z] - plrl.v[Z]) / dv;
   
   rhoc = get_element(&rho, &extent, nlc.v[X], nlc.v[Y], nlc.v[Z]);
   *rhoc += lllq;
   
    rhoc = get_element(&rho, &extent, nrc.v[X], nlc.v[Y], nlc.v[Z]);
   *rhoc += rllq;
   
   rhoc = get_element(&rho, &extent, nlc.v[X], nrc.v[Y], nlc.v[Z]);
   *rhoc += lrlq;
   
    rhoc = get_element(&rho, &extent, nlc.v[X], nlc.v[Y], nrc.v[Z]);
   *rhoc += llrq;
   
   rhoc = get_element(&rho, &extent, nrc.v[X], nrc.v[Y], nlc.v[Z]);
   *rhoc += rrlq;
   
   rhoc = get_element(&rho, &extent, nlc.v[X], nrc.v[Y], nrc.v[Z]);
   *rhoc += lrrq;
   
   rhoc = get_element(&rho, &extent, nrc.v[X], nlc.v[Y], nrc.v[Z]);
   *rhoc += rlrq;
   
   rhoc = get_element(&rho, &extent, nrc.v[X], nrc.v[Y], nrc.v[Z]);
   *rhoc += rrrq;
   __syncthreads();
}

void ucgrid3d_weight_particles_rho(particles_set* set, ucgrid3d* grid)
{
#ifdef VERBOSE
    printf("ucgrid3d_weight_particles_rho: weightening\n");
    printf("ucgrid3d_weight_particles_rho: using weightening method: %d\n",
           WEIGHTENING_METHOD, BOUNDARIES_TYPE);
#endif
#if ((WEIGHTENING_METHOD == WLINEAR) && (BOUNDARIES_TYPE == CYCLIC_BOUNDARIES))
    

    pos3d* ppos;
    speed3d* pspeed;
    particle* p = set->particles;
    int i;
    particle* pbunch = NULL;
    int nbunches = floor(set->num_of_particles/PBUNCHSIZE);
    cudaExtent extent = make_cudaExtent(sizeof(cufftComplex)*grid->Nx, grid->Ny, grid->Nz);
    #ifdef DEBUG
    printf("ucgrid3d_weight_particles_rho: total nbunches: %d\n", nbunches);
    #endif
    
    if (cudaMalloc(&pbunch, PBUNCHSIZE*sizeof(particle)) == cudaSuccess)
    {
        for (i = 0; i < nbunches; ++i)
        {
#ifdef DEBUG
            printf("ucgrid3d_weight_particles_rho: process %d bunch\n", i + 1);
#endif
            if (cudaMemcpy(pbunch, p, sizeof(particle)*PBUNCHSIZE, cudaMemcpyHostToDevice) == cudaSuccess)
            {
                do_weight_particles_rho<<<NUMOFBLOCKS, THREADS_PER_BLOCK>>>(grid->rho, extent, pbunch, set->charge, grid->dx, grid->dy, grid->dz);
                p = &(set->particles[i*PBUNCHSIZE]);
            }
#ifdef DEBUG
            else
            {
                printf("ucgrid3d_weight_particles_rho: cuda: memcpy error\n");
            }
#endif
        }
        cudaFree(pbunch);
    }
#ifdef DEBUG
    else
    {
        printf("ucgrid3d_weight_particles_rho: cuda: malloc error\n");
    }
#endif
    

/*
    floatdig dv = grid->dx * grid->dy * grid->dz;
    floatdig qdv = set->charge / dv;
    for (i = 0; i < set->num_of_particles; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            ppos = &(set->particles[i].pos);
            pspeed = &(set->particles[i].speed);
            ucgrid3d_get_neighbour_cells(grid, ppos, &nlc, &nrc);
            ucgrid3d_get_cell_pos(grid, &nlc, &plc);
            ucgrid3d_get_cell_pos(grid, &nrc, &prc);

            plll.v[X] = plc.v[X];
            plll.v[Y] = plc.v[Y];
             plll.v[Z] = plc.v[Z];

            prll.v[X] = prc.v[X];
            prll.v[Y] = plc.v[Y];
            prll.v[Z] = plc.v[Z];

            plrl.v[X] = plc.v[X];
            plrl.v[Y] = prc.v[Y];
            plrl.v[Z] = plc.v[Z];

            prrl.v[X] = prc.v[X];
            prrl.v[Y] = prc.v[Y];
            prrl.v[Z] = plc.v[Z];

            prrr.v[X] = prc.v[X];
            prrr.v[Y] = prc.v[Y];
            prrr.v[Z] = prc.v[Z];

            prlr.v[X] = prc.v[X];
            prlr.v[Y] = plc.v[Y];
            prlr.v[Z] = prc.v[Z];

            pllr.v[X] = plc.v[X];
            pllr.v[Y] = plc.v[Y];
            pllr.v[Z] = prc.v[Z];

            plrr.v[X] = plc.v[X];
            plrr.v[Y] = prc.v[Y];
            plrr.v[Z] = prc.v[Z];

            lllq = qdv * (prll.v[X] - ppos->v[X])*(plrl.v[Y] - ppos->v[Y])*(pllr.v[Z] - ppos->v[Z]) / dv;
            rllq = qdv * (ppos->v[X] - plll.v[X])*(prrl.v[Y] - ppos->v[Y])*(prlr.v[Z] - ppos->v[Z]) / dv;

            lrlq = qdv * (prrl.v[X] - ppos->v[X])*(ppos->v[Y] - plll.v[Y])*(plrr.v[Z] - ppos->v[Z]) / dv;
            rrlq = qdv * (ppos->v[X] - plrl.v[X])*(ppos->v[Y] - prll.v[Y])*(prrr.v[Z] - ppos->v[Z]) / dv;

            rrrq = qdv * (ppos->v[X] - plrr.v[X])*(ppos->v[Y] - prlr.v[Y])*(ppos->v[Z] - prrl.v[Z]) / dv;
            rlrq = qdv * (ppos->v[X] - pllr.v[X])*(prrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - prll.v[Z]) / dv;

            llrq = qdv * (prlr.v[X] - ppos->v[X])*(plrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - plll.v[Z]) / dv;
            lrrq = qdv * (prrr.v[X] - ppos->v[X])*(ppos->v[Y] - pllr.v[Y])*(ppos->v[Z] - plrl.v[Z]) / dv;

            grid->rho[nlc.v[X]][nlc.v[Y]][nlc.v[Z]] += lllq;

            grid->rho[nrc.v[X]][nlc.v[Y]][nlc.v[Z]] += rllq;

            grid->rho[nlc.v[X]][nrc.v[Y]][nlc.v[Z]] += lrlq;

            grid->rho[nlc.v[X]][nlc.v[Y]][nrc.v[Z]] += llrq;

            grid->rho[nrc.v[X]][nrc.v[Y]][nlc.v[Z]] += rrlq;

            grid->rho[nlc.v[X]][nrc.v[Y]][nrc.v[Z]] += lrrq;

            grid->rho[nrc.v[X]][nlc.v[Y]][nrc.v[Z]] += rlrq;

            grid->rho[nrc.v[X]][nrc.v[Y]][nrc.v[Z]] += rrrq;
        }
    }
*/
#endif
}

//void ucgrid3d_dump_rho_into_file(ucgrid3d* grid, char* homedir, char* filename)
//{
//    int i = 0;
//    int j = 0;
//    int k = 0;
//    char fullname[MAX_BUF_SIZE];
//    FILE* to;
//    sprintf(fullname, "%s/%s", homedir, filename);
//    to = fopen(fullname, "w");
//    fprintf(to, "%d\t%d\t%d\n", grid->Nx, grid->Ny, grid->Nz);
//    for (i = 0; i < grid->Nx; i++)
//    {
//        for (j = 0; j < grid->Ny; j++)
//        {
//            for (k = 0; k < grid->Nz; k++)
//            {
//                fprintf(to, "%e\n", grid->rho[i][j][k]);
//            }
//        }
//    }
//    fclose(to);
//}
